use crate::model::satori::SatoriMessage;
use crate::util::Config;
use axum::extract::ws::{Message, WebSocket};
use axum::http::HeaderMap;
use dashmap::DashMap;
use futures_util::stream::SplitSink;
use futures_util::StreamExt;
use sqlx::{Pool, Postgres};
use std::sync::Arc;
use tokio::sync::mpsc::Sender;
use uuid::Uuid;
use void_log::log_info;

mod receive;
mod send;

#[derive(Clone)]
pub struct AppState {
    pub clients: Arc<DashMap<Uuid, SplitSink<WebSocket, Message>>>,
    pub sender_tx: Sender<SenderData>,
    pub pool: Pool<Postgres>,
    pub user_token: Arc<DashMap<String, Uuid>>,
}

#[derive(Debug, Clone)]
pub struct SenderData {
    pub to_user: String,
    pub satori_message: SatoriMessage,
}

impl AppState {
    pub async fn new(sender_tx: Sender<SenderData>) -> Self {
        Self {
            clients: Arc::new(Default::default()),
            sender_tx,
            pool: Config::pg_conn().await,
            user_token: Arc::new(Default::default()),
        }
    }

    pub async fn handle_socket(self, socket: WebSocket, headers: HeaderMap) {
        // 分离发送线程与接受线程
        let (sender, receiver) = socket.split();

        if let Some(authorization) = headers.get("authorization") {
            // the userid is user_token
            let user_token = authorization.to_str().unwrap().trim_start_matches("Bearer ");

            log_info!("登录 {user_token}");

            // 查询userid是否登录
            if let Some(token) = self.user_token.get(user_token) {
                let userid = token.value().clone();
                // 保存握手信息
                self.clients.insert(userid, sender);

                log_info!("Client '{}' connected", &userid);

                //启用接收
                tokio::spawn(self.clone().receive_task(receiver, userid));
            }

            // 下线清除登录信息
            self.user_token.remove(user_token);
        };
    }
}